import asyncio
from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import JsonOutputParser
import os

key = os.getenv("OPENAI_API_KEY")
# print(key)
api_key = str(key)

model = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    base_url="https://api.zetatechs.com/v1",
    api_key=api_key
)

# events = []
# async def main():
#     async for event in model.astream_events("hello"):
#         events.append(event)
#
# asyncio.run(main())
# 这里的 *是一个解包操作, 将events列表中的元素作为独立的参数传递给print函数
# print(*events[:3],sep="\n")
# print(*events[-2:],sep="\n")

# chain 式的Stream events
chain = (
        model | JsonOutputParser()
)  # Due to a bug in older versions of Langchain, JsonOutputParser did not stream results from some models


async def use():
    events = [
        event
        async for event in chain.astream_events(
            "output a list of the countries france, spain and japan and their populations in JSON format. "
            'Use a dict with an outer key of "countries" which contains a list of countries. '
            "Each country should have the key `name` and `population`",
        )
    ]
    print(events[:3], sep='\n')


# asyncio.run(use())


async def use1():
    num_events = 0
    async for event in chain.astream_events(
            "output a list of the countries france, spain and japan and their populations in JSON format. "
            'Use a dict with an outer key of "countries" which contains a list of countries. '
            "Each country should have the key `name` and `population`",
    ):
        kind = event["event"]
        if kind == "on_chat_model_stream":
            print(
                f"Chat model chunk: {repr(event['data']['chunk'].content)}",
                flush=True,
            )
        if kind == "on_parser_stream":
            print(f"Parser chunk: {event['data']['chunk']}", flush=True)
        num_events += 1
        if num_events > 30:
            # Truncate the output
            print("...")
            break


asyncio.run(use1())